'''
自监督攻击, 攻击干净样本生成对抗样本用于训练净化网络purify
Self supervised attack is purely based on maximizing the perceptual feature difference.最大化感知特征差异
'''
# 重写图像保存的实现

import os
import sys
sys.path.append('../')
import argparse
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.utils import save_image
import torchvision
import numpy as np
import random
from utils.PairedImageDataset_util import AdvImageDataset

# 通过设置随机种子可使训练结果可复现。
def set_random_seed(seed):
    """Set random seeds."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

#################################################################
# 特征提取
# VGG [42] network pretrained on ImageNet.
# 浅层特征抽取网络，提取浅层特征。
# 低分辨率图像进入后会经过一个卷积+RELU函数，将输入通道数调整为64??
class FeatureExtractor(nn.Module): # perceptual_criteria 感知标准，继承nn.Module类
    def __init__(self, ssp_layer):               # 类的构造函数或初始化方法
        super(FeatureExtractor, self).__init__() 
        # you can try other models
        alexnet = torchvision.models.alexnet(pretrained=True) 
        self.feature_extractor = nn.Sequential(*list(alexnet.features))[:ssp_layer].eval() 
    def forward(self, img): 
        return self.feature_extractor(img)

# argparse是一个Python模块：命令行选项、参数和子命令解析器。
parser = argparse.ArgumentParser(description='SSP Attack') # 创建解析器
parser.add_argument('--sourcedir', default='Data/NIPS2017/images') # default='clean_imgs'
parser.add_argument('--targetdir', default='Data/alexnet/adv_SSP_alexnet_conv1_eps16_iters100_MSE_normTrue')
parser.add_argument('--batch_size', type=int, default=64, help='Batch size')
parser.add_argument('--eps', type=int, default= 16,  help ='perturbation budget')
parser.add_argument('--step_size', type=float, default=0.01, help='Step size')
parser.add_argument('--iters', type=int, default=100, help='Number of SSP Iterations')
parser.add_argument('--ssp_layer', type=int, default=2, help='Model layer that is going to be used in SSP')

args = parser.parse_args() # parse_args(args=None, nampespace=None)parse_args(),将之前add_argument()定义的参数进行赋值，并返回相关的namespace
print(args)

set_random_seed(0) # 设置随机种子

# GPU or CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# device = torch.device('cpu')
# 此后的criterion.to(device)，img = img.to(device)等都是将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去，之后的运算都在GPU上进行。

os.makedirs(args.targetdir, exist_ok=True)

feature_extractor = FeatureExtractor(args.ssp_layer).to(device)
criterion = torch.nn.MSELoss().to(device)
eps = args.eps / 255 # 扰动预算，现在设置的是16，16 / 255
step = args.step_size # 步长

trans = transforms.Compose([
        transforms.Resize(256),       # 256, 将图片短边缩放至x = 256, 长宽比保持不变; transforms.Resize([224, 224])就能将输入图片转化成224×224的输入特征图
        transforms.CenterCrop(224),   # 224, 功能：从图像中心裁剪图片; size：所需裁剪图片尺寸; 图片从中心点开始计算, 左右宽224, 上下高224的一个裁剪区域
        transforms.ToTensor()
    ])
dataset = AdvImageDataset(img_dirs=args.sourcedir,
                        transform=trans)
data_loader = torch.utils.data.DataLoader(
    dataset,                     # 数据加载器，将数据放在数据库中
    batch_size=args.batch_size,  # 从数据库中每次抽出batch size个样本
    shuffle=False, 
    num_workers=4,               # 用于数据加载的子进程数量，0意味着数据将在主进程中加载
    pin_memory=True)             # 数据加载器将在返回张量之前，将张量复制到CUDA固定内存中

# 归一化函数
mean = [0.485, 0.456, 0.406] # 数据集的均值
std = [0.229, 0.224, 0.225] # 数据集的方差
# 对数据进行归一化的操作 image = (image - mean) / std
def normalize(t):
    t[:, 0, :, :] = (t[:, 0, :, :] - mean[0])/std[0]
    t[:, 1, :, :] = (t[:, 1, :, :] - mean[1])/std[1]
    t[:, 2, :, :] = (t[:, 2, :, :] - mean[2])/std[2]
    return t

for i, (img, img_name) in enumerate(data_loader):
    # Generating adversarial image
    img = img.to(device)
    # adv初始化，与干净样本同样大小
    adv = torch.randn(img.shape).to(device) # img.shape返回的是三维数组([0]high,[1]width,[2]3)；而torch.randn()返回一个张量，包含了从标准正态分布（均值为0，方差为1，即高斯白噪声）中抽取的一组随机数。
    adv.requires_grad = True # 设置 require_grad 为 True 意为梯度反传时对该 Tensor 计算梯度，并存入 tensor.grad 中。即adv.grad

    # for t in range(100): # iterations迭代T，range(start, stop[, step])，start默认为0，故range(100)等价于range[0,100)，从0到100，不包含100；step步长，默认为1
    for t in range(args.iters):
        adv1 = adv + 0 # 用于归一化计算loss，而不改变adv
        loss = criterion(feature_extractor(normalize(adv1)), feature_extractor(normalize(img.clone().detach())))
        loss.backward() # 反向传播

        adv.data = adv.data + step * adv.grad.sign() # Generate adversaries
        adv.data = torch.where(adv.data > img.data + eps, img.data + eps, adv.data)
        adv.data = torch.where(adv.data < img.data - eps, img.data - eps, adv.data)
        adv.data.clamp_(0.0, 1.0) # torch.clamp(input, min, max, out=None) → Tensor。input (Tensor) – 输入张量，min (Number) – 限制范围下限，max (Number) – 限制范围上限，out (Tensor, optional) – 输出张量
        adv.grad.data.zero_() # 在grad更新时，每一次运算后都需要将上一次的梯度记录清空

    for img_index in range(img.size()[0]):
        adv_name = img_name[img_index]
        adv_path = os.path.join(args.targetdir, adv_name)
        save_image(adv[img_index], adv_path, nrow=1, normalize=False)
    
    print('Number of processed images:', args.batch_size*(i+1))